All elements in two binary search tree [Stack]

Time: O(N); Space: O(H); medium

Given two binary search trees root1 and root2.

Return a list containing all the integers from both trees sorted in ascending order.

Example 1:

Input: root1 = {TreeNode} [2,1,4], root2 = {TreeNode} [1,0,3]

Output: [0,1,1,2,3,4]

Example 2:

Input: root1 = {TreeNode} [0,-10,10], root2 = {TreeNode} [5,1,7,0,2]

Output: [-10,0,0,1,2,5,7,10]

Example 3:

Input: root1 = {TreeNode} [], root2 = {TreeNode} [5,1,7,0,2]

Output: [0,1,2,5,7]

Example 4:

Input: root1 = {TreeNode} [0,-10,10], root2 = {TreeNode} []

Output: [-10,0,10]

Example 5:

Input: root1 = {TreeNode} [1,null,8], root2 = {TreeNode} [8,1]

Output: [1,1,8,8]

Constraints:

  • Each tree has at most 5000 nodes.

  • Each node’s value is between [-10^5, 10^5].

Hints:

  1. Traverse the first tree in list1 and the second tree in list2.

  2. Merge the two trees in one list and sort it.

[4]:
class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
[5]:
class Solution1(object):
    """
    Time: O(N)
    Space: O(H)
    """
    def getAllElements(self, root1, root2):
        """
        :type root1: TreeNode
        :type root2: TreeNode
        :rtype: List[int]
        """
        def inorder_gen(root):
            result, stack = [], [(root, False)]
            while stack:
                root, is_visited = stack.pop()
                if root is None:
                    continue
                if is_visited:
                    yield root.val
                else:
                    stack.append((root.right, False))
                    stack.append((root, True))
                    stack.append((root.left, False))
            yield None

        result = []
        left_gen, right_gen = inorder_gen(root1), inorder_gen(root2)
        left, right = next(left_gen), next(right_gen)

        while left is not None or right is not None:
            if right is None or (left is not None and left < right):
                result.append(left)
                left = next(left_gen)
            else:
                result.append(right)
                right = next(right_gen)
        return result
[8]:
s = Solution1()

root1 = TreeNode(2)
root1.left = TreeNode(1)
root1.right = TreeNode(4)
root2 = TreeNode(1)
root2.left = TreeNode(0)
root2.right = TreeNode(3)
assert s.getAllElements(root1, root2) == [0,1,1,2,3,4]

root1 = TreeNode(0)
root1.left = TreeNode(-10)
root1.right = TreeNode(10)
root2 = TreeNode(5)
root2.left = TreeNode(1)
root2.right = TreeNode(7)
root2.left.left = TreeNode(0)
root2.left.right = TreeNode(2)
assert s.getAllElements(root1, root2) == [-10,0,0,1,2,5,7,10]

root1 = None
root2 = TreeNode(5)
root2.left = TreeNode(1)
root2.right = TreeNode(7)
root2.left.left = TreeNode(0)
root2.left.right = TreeNode(2)
assert s.getAllElements(root1, root2) == [0,1,2,5,7]

root1 = TreeNode(0)
root1.left = TreeNode(-10)
root1.right = TreeNode(10)
root2 = None
assert s.getAllElements(root1, root2) == [-10,0,10]

root1 = TreeNode(1)
root1.right = TreeNode(8)
root2 = TreeNode(8)
root2.left = TreeNode(1)
assert s.getAllElements(root1, root2) == [1,1,8,8]